import os
import numpy as np
from ppgnss import gnss_io
from ppgnss import gnss_utils
import matplotlib
import matplotlib.pyplot as plt

fig_dir = "../figures"
sites = ["B216", "B240", "B286", "PD02"]
doys = [98, 99, 100, 101, 102, 103, 104]
for doy in doys:
    for site in sites:
        fig = plt.figure(figsize=(80, 80))
        fig, axes = plt.subplots( nrows=3, ncols=3, subplot_kw={'projection': 'polar'})
        xtr_filename = "../data/out/%s%03d.xtr" % (site, doy)
        obj_filename = "../data/snr/%s%03d.obj" % (site, doy)
        if os.path.isfile(obj_filename):
            data = gnss_utils.loadobject(obj_filename)
            print(obj_filename)
            xr_eles, xr_azis, xr_snrs = data["ele"], data["azi"], data["snr"]
        else:
            xr_eles, xr_azis, xr_snrs = gnss_io.read_anubis(xtr_filename)
            print(xtr_filename)
            data = {"ele": xr_eles,
                    "azi": xr_azis,
                    "snr": xr_snrs}
            gnss_utils.saveobject(data, obj_filename)
        for iband, band in enumerate(xr_snrs.coords["band"]):
            irow = iband//3
            icol = iband%3
            ax = axes[irow][icol]
            strband = str(band.values)
            sattype = strband[:3]
            xr_ele = xr_eles.loc[sattype]
            xr_azi = xr_azis.loc[sattype]
            xr_snr = xr_snrs.loc[strband]
            idx = xr_ele.values > 0 
            ele = xr_ele.values[idx]
            azi = xr_azi.values[idx]
            snr = xr_snr.values[idx]
            im = ax.scatter(np.deg2rad(azi), np.pi/2-np.deg2rad(ele), s=2, c=snr, cmap=matplotlib.colormaps["RdYlBu"], vmin=20, vmax=70)
            ax.set_theta_zero_location("N")
            ax.set_theta_direction(-1)
            ax.set_title(strband)
        plt.colorbar(im)
        fig_filename = os.path.join(fig_dir, "%s_%04d_%03d.png" %(site, 2022, doy))
        plt.savefig(fig_filename, dpi=600)
        plt.close()
        